import json
from typing import Tuple
from app.config import GremlinGdbConfig
from gremlin_python.driver.driver_remote_connection import DriverRemoteConnection
from gremlin_python.process.anonymous_traversal import traversal
from gremlin_python.process.graph_traversal import GraphTraversalSource


class InitGremlinGdbSchemeForPromqt:
    """ Init Gremlin Gdb Scheme for Promqt """

    def __init__(self, url: str = None, password: str = None, username: str = None) -> None:
        """ Init InitGremlinGdbSchemeForPromqt properties
        :param url: gremlin gdb url
        :param password: gremlin gdb password
        :param username: gremlin gdb username
        """

        self.obj = GremlinGdbConfig()
        self.url = url or self.obj.gremlin_url
        self.password = password or self.obj.gremlin_password
        self.username = username or self.obj.gremlin_username

    def init_gremlin_g_handler(self) -> Tuple[DriverRemoteConnection, GraphTraversalSource]:
        """ init Gremlin Gdb Connection and 'g' handler """

        connection = DriverRemoteConnection(url=self.url, username=self.username, password=self.password)
        g = traversal().with_remote(connection)
        return connection, g.with_('scriptEvaluationTimeout', 60000)

    def get_all_labels_count(self) -> int:
        """ get all labels count """

        conn, g = self.init_gremlin_g_handler()
        count = g.V().count().next()
        conn.close()
        return count

    def get_schemes_labels(self, batch_size: int = 1000, offset: int = 0) -> list:
        """ get gremlin gdb all V Labels """

        conn, g = self.init_gremlin_g_handler()
        r = list(dict(g.V().label().groupCount().next()).keys())
        conn.close()
        return r

    def get_schemes_labels_threading(self, batch_size: int = 50000, num_threads: int = 10):
        """ Aborted! """

        import threading
        from queue import Queue
        from tqdm import tqdm

        conn, g = self.init_gremlin_g_handler()
        global result
        result = []
        global error_result
        error_result = []

        total_count = self.get_all_labels_count()
        task_queue = Queue()
        lock = threading.Lock()
        pbar = tqdm(total=total_count, desc="Processing", unit="labels")

        for start_offset in range(0, total_count, batch_size):
            end_offset = min(start_offset + batch_size, total_count)
            task_queue.put((start_offset, end_offset))

        def worker(queue: Queue, g: GraphTraversalSource, lock: threading.Lock, pbar: tqdm):
            while True:
                try:
                    start_offset, end_offset = queue.get(block=False)
                except queue.Empty:
                    break
                except Exception as e:
                    task_queue.put((start_offset, end_offset))
                    break
                labels = g.V().range_(start_offset, end_offset).label().dedup().toList()
                with lock:
                    result.extend(labels)
                    pbar.update(end_offset - start_offset)
                queue.task_done()

        threads = []
        for _ in range(num_threads):
            t = threading.Thread(target=worker, args=(task_queue, g, lock, pbar))
            t.start()
            threads.append(t)

        task_queue.join()

        for t in threads:
            t.join()

        pbar.close()
        with open('result.json', 'w') as f:
            f.write(json.dumps(result, ensure_ascii=False))
        with open('err_result.json', 'w') as f:
            f.write(json.dumps(error_result, ensure_ascii=False))
        return result

    def get_label_with_properties(self, labels: list, need_format: bool = False) -> list:
        """ get all label with properties """

        conn, g = self.init_gremlin_g_handler()
        res = []
        for label in labels:
            v = g.V().hasLabel(label).limit(1).valueMap().next()
            res.append({
                'label': label,
                'properties': list(v.keys()),
            })

        conn.close()
        return res if not need_format else json.dumps(res, ensure_ascii=False, indent=4)

    def execute_data_result(self):

        conn, g = self.init_gremlin_g_handler()
        res = g.V().hasLabel('ib_job').has('jobName', '程序员').values('aiAbilityList').toList()
        conn.close()
        return res


if __name__ == '__main__':
    """
        我学什么技能 才能做程序员 
        我想做程序员 我该学什么技能
    """
    obj = InitGremlinGdbSchemeForPromqt()
    # all_labels = obj.get_schemes_labels()
    # label_with_properties = obj.get_label_with_properties(all_labels)
    # print(json.dumps(label_with_properties, ensure_ascii=False, indent=4))
    print(obj.execute_data_result())